import os
import numpy as np
import json
import torch
import torch.nn as nn
from meshsegnet import *
import utils
import vedo
import pandas as pd
import operator
from losses_and_metrics_for_mesh import *
from scipy.spatial import distance_matrix

from utils import color_2_label, label_2_color

os.environ["KMP_DUPLICATE_LIB_OK"]="TRUE"

if __name__ == '__main__':
    torch.cuda.set_device(utils.get_avail_gpu())

    with open('config.json', 'r') as f:
        config = json.load(f)

    mod = config["mod"]
    test_config = config["test"]

    model_path = config["train"]["model_path"]
    model_name = '{}{}_best.tar'.format(config["train"]["model_name"], mod)

    mesh_path = config["data"]["data_path"] + mod
    label_path = config["data"]["label_data_path"] + mod

    test_list = pd.read_csv(config["list"]["list_path"] + 'test_list_1.csv')['Test ID'].values
    test_mesh_filename = '{}{}.ply'
    test_path = test_config["test_path"]

    if not os.path.exists(test_path):
        os.mkdir(test_path)

    num_classes = test_config["num_classes"]
    num_channels = test_config["num_channels"]

    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    model = MeshSegNet(num_classes=num_classes, num_channels=num_channels).to(device, dtype=torch.float)

    checkpoint = torch.load(os.path.join(model_path, model_name), map_location='cpu')
    model.load_state_dict(checkpoint['model_state_dict'])

    torch.backends.cudnn.benchmark = True
    torch.backends.cudnn.enabled = True

    dsc = []
    sen = []
    ppv = []

    print('Testing')
    model.eval()
    with torch.no_grad():
        for i_sample in test_list:

            print('Predicting Sample filename: {}'.format(test_mesh_filename.format(mod, i_sample)))

            mesh = vedo.load(os.path.join(mesh_path, test_mesh_filename.format(mod, i_sample)))
            label_mesh = vedo.load(os.path.join(label_path, test_mesh_filename.format(mod, i_sample)))

            colors = mesh.cellcolors[:, 0:3] / 255.
            labels = color_2_label(label_mesh.cellcolors, mod=mod)
            predicted_labels = np.zeros(labels.shape)

            points = mesh.points()
            mean_cell_centers = mesh.center_of_mass()
            points[:, 0:3] -= mean_cell_centers[0:3]

            ids = np.array(mesh.faces())
            cells = points[ids].reshape(mesh.ncells, 9).astype(dtype='float32')

            mesh.compute_normals()
            normals = mesh.celldata['Normals']

            barycenters = mesh.cell_centers()
            barycenters -= mean_cell_centers[0:3]

            maxs = points.max(axis=0)
            mins = points.min(axis=0)
            means = points.mean(axis=0)
            stds = points.std(axis=0)
            nmeans = normals.mean(axis=0)
            nstds = normals.std(axis=0)

            for i in range(3):
                cells[:, i] = (cells[:, i] - means[i]) / stds[i]
                cells[:, i+3] = (cells[:, i+3] - means[i]) / stds[i]
                cells[:, i+6] = (cells[:, i+6] - means[i]) / stds[i]
                barycenters[:, i] = (barycenters[:, i] - mins[i]) / (maxs[i] - mins[i])
                normals[:, i] = (normals[:, i] - nmeans[i]) / nstds[i]

            X = np.column_stack((cells, barycenters, normals, colors))
            Y = labels



            # # valid_idx = np.argwhere(labels >= 0)[:, 0]
            # #
            # # selected_idx = np.random.choice(valid_idx, size=6000, replace=False)
            # # selected_idx = np.sort(selected_idx, axis=None)
            #
            # positive_idx = np.argwhere(labels > 0)[:, 0]  # tooth idx
            # negative_idx = np.argwhere(labels == 0)[:, 0]  # gingiva idx
            #
            # num_positive = len(positive_idx)  # number of selected tooth cells
            #
            # X_test = np.zeros([num_positive, X.shape[1]], dtype='float32')
            # Y_test = np.zeros([num_positive, Y.shape[1]], dtype='int32')
            # S1 = np.zeros([num_positive, num_positive], dtype='float32')
            # S2 = np.zeros([num_positive, num_positive], dtype='float32')
            #
            # # if num_positive > 6000:  # all positive_idx in this patch
            # #     positive_selected_idx = np.random.choice(positive_idx, size=6000, replace=False)
            # #     selected_idx = positive_selected_idx
            # # else:  # patch contains all positive_idx and some negative_idx
            # #     num_negative = 6000 - num_positive  # number of selected gingiva cells
            # #     positive_selected_idx = np.random.choice(positive_idx, size=num_positive, replace=False)
            # #     negative_selected_idx = np.random.choice(negative_idx, size=num_negative, replace=False)
            # #     selected_idx = np.concatenate((positive_selected_idx, negative_selected_idx))
            # #
            # # selected_idx = np.sort(selected_idx, axis=None)
            # selected_idx = np.random.choice(positive_idx, size=num_positive, replace=False)
            #
            # X_test[:] = X[selected_idx, :]
            # Y_test[:] = Y[selected_idx, :]
            #
            # # print(Y_test.shape, labels.shape)
            #
            # predicted_labels = np.zeros(Y_test.shape)
            #
            # if torch.cuda.is_available():
            #     TX = torch.as_tensor(X_test[:, 9:12], device='cuda')
            #     TD = torch.cdist(TX, TX)
            #     D = TD.cpu().numpy()
            # else:
            #     D = distance_matrix(X_test[:, 9:12], X_test[:, 9:12])
            #
            # S1[D < 0.1] = 1.0
            # S1 = S1 / np.dot(np.sum(S1, axis=1, keepdims=True), np.ones((1, num_positive)))
            #
            # S2[D < 0.2] = 1.0
            # S2 = S2 / np.dot(np.sum(S2, axis=1, keepdims=True), np.ones((1, num_positive)))
            #
            # X_test = X_test.transpose(1, 0)
            # # Y_test = Y_test.transpose(1, 0)
            #
            # X_test = X_test.reshape([1, X_test.shape[0], X_test.shape[1]])
            # S1 = S1.reshape([1, S1.shape[0], S1.shape[1]])
            # S2 = S2.reshape([1, S2.shape[0], S2.shape[1]])
            #
            # X_test = torch.from_numpy(X_test).to(device, dtype=torch.float)
            # S1 = torch.from_numpy(S1).to(device, dtype=torch.float)
            # S2 = torch.from_numpy(S2).to(device, dtype=torch.float)
            #
            # tensor_prob_output = model(X_test, S1, S2).to(device, dtype=torch.float)
            # patch_prob_output = tensor_prob_output.cpu().numpy()
            #
            # for i_label in range(num_classes):
            #     predicted_labels[np.argmax(patch_prob_output[0, :], axis=-1) == i_label] = i_label

            # computing A_S and A_L
            A_S = np.zeros([X.shape[0], X.shape[0]], dtype='float32')
            A_L = np.zeros([X.shape[0], X.shape[0]], dtype='float32')
            D = distance_matrix(X[:, 9:12], X[:, 9:12])
            A_S[D < 0.1] = 1.0
            A_S = A_S / np.dot(np.sum(A_S, axis=1, keepdims=True), np.ones((1, X.shape[0])))

            A_L[D < 0.2] = 1.0
            A_L = A_L / np.dot(np.sum(A_L, axis=1, keepdims=True), np.ones((1, X.shape[0])))

            # numpy -> torch.tensor
            X = X.transpose(1, 0)
            X = X.reshape([1, X.shape[0], X.shape[1]])
            X = torch.from_numpy(X).to(device, dtype=torch.float)
            A_S = A_S.reshape([1, A_S.shape[0], A_S.shape[1]])
            A_L = A_L.reshape([1, A_L.shape[0], A_L.shape[1]])
            A_S = torch.from_numpy(A_S).to(device, dtype=torch.float)
            A_L = torch.from_numpy(A_L).to(device, dtype=torch.float)

            tensor_prob_output = model(X, A_S, A_L).to(device, dtype=torch.float)
            patch_prob_output = tensor_prob_output.cpu().numpy()

            for i_label in range(num_classes):
                predicted_labels[np.argmax(patch_prob_output[0, :], axis=-1) == i_label] = i_label

            # output predicted labels
            mesh2 = mesh.clone()
            mesh2.cellcolors = label_2_color(predicted_labels, mod, mesh2.cellcolors, selected_idx=None)
            vedo.write(mesh2, os.path.join(test_path, '{}{}.ply'.format(mod, i_sample)), binary=False)

            # print(predicted_labels.shape, Y_test.shape)
            # convert predict result and label to one-hot maps
            tensor_predicted_labels = torch.from_numpy(predicted_labels)
            tensor_test_labels = torch.from_numpy(labels)
            tensor_predicted_labels = tensor_predicted_labels.long()
            tensor_test_labels = tensor_test_labels.long()
            one_hot_predicted_labels = nn.functional.one_hot(tensor_predicted_labels[:, 0], num_classes=num_classes)
            one_hot_labels = nn.functional.one_hot(tensor_test_labels[:, 0], num_classes=num_classes)

            # print(Y_test)
            # print(one_hot_labels)
            # print(one_hot_predicted_labels.shape, one_hot_labels.shape)
            print(one_hot_predicted_labels.sum(dim=0))
            print(one_hot_labels.sum(dim=0))

            # calculate DSC
            i_dsc = DSC(one_hot_predicted_labels, one_hot_labels)
            i_sen = SEN(one_hot_predicted_labels, one_hot_labels)
            i_ppv = PPV(one_hot_predicted_labels, one_hot_labels)
            dsc.append(i_dsc)
            sen.append(i_sen)
            ppv.append(i_ppv)
            # print('\tLabel 1: {}; Label 2: {}'.format(dsc[0], dsc[1]))

        dsc = np.asarray(dsc)
        sen = np.asarray(sen)
        ppv = np.asarray(ppv)

        # output all DSCs
        all_dsc = pd.DataFrame(data=dsc, index=test_list, columns=['label {}'.format(i) for i in range(1, num_classes)])
        all_sen = pd.DataFrame(data=sen, index=test_list, columns=['label {}'.format(i) for i in range(1, num_classes)])
        all_ppv = pd.DataFrame(data=ppv, index=test_list, columns=['label {}'.format(i) for i in range(1, num_classes)])
        print(all_dsc)
        print(all_dsc.describe())
        print(all_sen)
        print(all_sen.describe())
        print(all_ppv)
        print(all_ppv.describe())
        all_dsc.to_csv(os.path.join(test_path, 'test_DSC_report_fold_1.csv'), header=True, index=True)
        all_sen.to_csv(os.path.join(test_path, 'test_SEN_report_fold_1.csv'), header=True, index=True)
        all_ppv.to_csv(os.path.join(test_path, 'test_PPV_report_fold_1.csv'), header=True, index=True)